import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from skimage.metrics import structural_similarity as ssim
import cv2
import os
import pytorch_fid.fid_score as fid_score
import torch
from torch_fidelity import calculate_metrics
from skimage.io import imread
from skimage.transform import resize
from skimage.metrics import peak_signal_noise_ratio
from inception_score_pytorch import inception_score as IS
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import torchvision


class CustomImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_files = [os.path.join(root_dir, file) for file in os.listdir(root_dir)]

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        img_path = self.image_files[idx]
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        if self.transform:
            image = self.transform(image)
        return image


def calculate_inception_score(generated_images, batch_size=32):
    # 定义图像变换
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将 PIL 图像转换为张量，范围 [0, 1]
        transforms.Resize((224, 224)),  # 可根据需要调整图像大小
    ])
    # 初始化数据集
    dataset = CustomImageDataset(generated_images, transform=transform)
    mean, std = IS.inception_score(dataset, cuda=True, batch_size=batch_size, resize=True, splits=1)
    return mean, std


def calculate_fid(generated_images, real_images):
    # 计算真实图像的数量
    num_real_images = len(os.listdir(real_images))
    # 计算生成图像的数量
    num_generated_images = len(os.listdir(generated_images))

    # 将 batch_size 设置为两者中较小的数量
    batch_size = min(num_real_images, num_generated_images)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    fid_value = fid_score.calculate_fid_given_paths([real_images, generated_images],
                                                    batch_size,
                                                    device,
                                                    2048)
    return fid_value


def calculate_psnr(generated_image, original_image):
    if original_image.shape != generated_image.shape:
        generated_image = resize(generated_image, original_image.shape)

    psnr_value = peak_signal_noise_ratio(original_image, generated_image)

    return psnr_value


def calculate_ssim(generated_image, original_image):
    if original_image.shape != generated_image.shape:
        generated_image = resize(generated_image, original_image.shape)

    # 检查图像的数据类型，如果是浮点数，设置 data_range 参数
    if original_image.dtype == np.float32 or original_image.dtype == np.float64:
        data_range = 1.0  # 假设图像的像素值范围是 [0, 1]
    else:
        data_range = 255.0  # 假设图像的像素值范围是 [0, 255]

    # 计算 SSIM
    ssim_value = ssim(original_image, generated_image, multichannel=True, channel_axis=2, data_range=data_range)

    return ssim_value


def load_images_from_dir(image_dir):
    images = []
    target_size = (512, 512)
    for file in os.listdir(image_dir):
        img_path = os.path.join(image_dir, file)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # 调整图像大小
        img = cv2.resize(img, target_size)
        img = img.astype(np.float32) / 255.0
        images.append(img)
    return np.array(images)


def evaluate_images(generated_image_dir, original_image_dir):
    # 读取原始图像
    original_images = load_images_from_dir(original_image_dir)
    # 读取生成图像
    generated_images = load_images_from_dir(generated_image_dir)

    # 计算指标
    # mean_is, std_is = calculate_inception_score(generated_image_dir, batch_size=8)
    # fid_score = calculate_fid(generated_image_dir, original_image_dir)
    psnr_scores = []
    ssim_scores = []
    for gen_img, orig_img in zip(generated_images, original_images):
        psnr_scores.append(calculate_psnr(gen_img, orig_img))
        ssim_scores.append(calculate_ssim(gen_img, orig_img))
    avg_psnr_score = np.mean(psnr_scores)
    avg_ssim_score = np.mean(ssim_scores)
    return 0, 0, 0, avg_psnr_score, avg_ssim_score


if __name__ == "__main__":
    # dall_image_dir = "./dall_e3"1
    # generated_image_dir = "./outputs"
    # original_image_dir = "./truths"
    # mean_is, std_is, fid_score, avg_psnr_score, avg_ssim_score = evaluate_images(generated_image_dir, original_image_dir)
    # print("Ours:")
    # print(f"Inception Score: {is_score}")
    # print(f"FID Score: {fid_score}")
    # print(f"Average PSNR Score: {avg_psnr_score}")
    # print(f"Average SSIM Score: {avg_ssim_score}")
    # mean_is, std_is, fid_score, avg_psnr_score, avg_ssim_score = evaluate_images(dall_image_dir, original_image_dir)
    # print("DALL-E3:")
    # print(f"Inception Score: {is_score}")
    # print(f"FID Score: {fid_score}")
    # print(f"Average PSNR Score: {avg_psnr_score}")
    # print(f"Average SSIM Score: {avg_ssim_score}")
    # 评估结果
    dall_image_dir = "/root/data/dall_e3_recent"
    generated_image_dir = "/root/data/outputs_recent"
    original_image_dir = "/root/data/truths_recent"
    mean_is, std_is, fid_score, avg_psnr_score, avg_ssim_score = evaluate_images(generated_image_dir, original_image_dir)
    print("Ours:")
    # print(f"Inception Score: {is_score}")
    # print(f"FID Score: {fid_score}")
    print(f"Average PSNR Score: {avg_psnr_score}")
    print(f"Average SSIM Score: {avg_ssim_score}")
    mean_is, std_is, fid_score, avg_psnr_score, avg_ssim_score = evaluate_images(dall_image_dir, original_image_dir)
    print("DALL-E3:")
    # print(f"Inception Score: {is_score}")
    # print(f"FID Score: {fid_score}")
    print(f"Average PSNR Score: {avg_psnr_score}")
    print(f"Average SSIM Score: {avg_ssim_score}")
